import sympy
import sys

def generate_shake_series(natoms, constraints, name=""):
  r = [[sympy.symbols('r_{%d%d}' % (i, j)) for j in range(natoms)] for i in range(natoms)]
  d = [[sympy.symbols('d_{%d%d}' % (i, j)) for j in range(natoms)] for i in range(natoms)]
  ruc = [[sympy.symbols('r^{uc}_{%d%d}' % (i, j)) for j in range(natoms)] for i in range(natoms)]
  m = [sympy.symbols("m%s" % (i)) for i in range(natoms)]
  l = [sympy.symbols("l%s" % (i)) for i in range(len(constraints))]
  A = sympy.Matrix([[0 for i in range(len(constraints))] for j in range(len(constraints))])
  dt = sympy.symbols('\\varDelta{t}')
  for ik, [k1, k2] in enumerate(constraints):
    for ikk, [kk1, kk2] in enumerate(constraints):
      elem = 0
      delta11 = int(k1 == kk1)
      delta12 = int(k1 == kk2)
      delta21 = int(k2 == kk1)
      delta22 = int(k2 == kk2)
      elem = ((delta11 - delta12) / m[k1] + (delta22 - delta21) / m[k2]) * r[kk1][kk2] * ruc[k1][k2]
      A[ik, ikk] = elem
  c = sympy.Matrix([[0] for i in range(len(constraints))])
  for ik, [k1, k2] in enumerate(constraints):
    c[ik, 0] = (ruc[k1][k2] ** 2 - d[k1][k2] ** 2) / (4 * dt ** 2)
  print('\\begin{equation}')
  print('A_{%s}=' % name)
  print(sympy.latex(A))
  print('\\end{equation}')
  print('\\begin{equation}')
  print('c_{%s}=' % name)
  print(sympy.latex(c))
  print('\\end{equation}')
  print('\\begin{align}')
  for ik, [k1, k2] in enumerate(constraints):
    A = sympy.Matrix([[0 for i in range(len(constraints))] for j in range(len(constraints))])
    quad = 0
    for ikk, [ikk1, ikk2] in enumerate(constraints):
      for jkk, [jkk1, jkk2] in enumerate(constraints):
        deltai11 = int(k1 == ikk1)
        deltai12 = int(k1 == ikk2)
        deltai21 = int(k2 == ikk1)
        deltai22 = int(k2 == ikk2)
        deltaj11 = int(k1 == jkk1)
        deltaj12 = int(k1 == jkk2)
        deltaj21 = int(k2 == jkk1)
        deltaj22 = int(k2 == jkk2)
        quad += ((deltai11 - deltai12)/m[k1] + (deltai22 - deltai21)/m[k2])*((deltaj11 - deltaj12)/m[k1] + (deltaj22 - deltaj21)/m[k2]) * r[ikk1][ikk2] * r[jkk1][jkk2] * l[ikk] * l[jkk]

    print('q_%d &=' % ik)
    print(sympy.latex(sympy.simplify(quad)))
    if ik < len(constraints) - 1:
      print('\\\\')
  print('\\end{align}')
print('\\section{mSHAKE of CH}')
generate_shake_series(2, [[0, 1]], "CH")
print('\\section{mSHAKE of CH2}')
generate_shake_series(3, [[0, 1], [0, 2]], "CH2")
print('\\section{mSHAKE of OH2}')
generate_shake_series(3, [[0, 1], [0, 2], [1, 2]], "OH2")
print('\\section{mSHAKE of CH3}')
generate_shake_series(4, [[0, 1], [0, 2], [0, 3]], "CH3")
